#include "typechecker.h"
#include <boost/foreach.hpp>

#define foreach BOOST_FOREACH

TypeChecker::TypeChecker(shared_ptr<Program> program) : program(program) { }

TypeChecker::~TypeChecker() { }

void TypeChecker::check() {
	TableGenerator generator(program);
	generator.generate();
	classTable = generator.getClassTable();
	methods = generator.getMethodsMap();
	
	if(classTable->hasErrors()) {
		foreach(std::string error, classTable->getErrors()) {
			errors.push_back(error);
		}
	}

	std::pair<std::string, shared_ptr<SymbolTable> > methodPair;
	foreach(methodPair, methods) {
		shared_ptr<SymbolTable> methodTable = methodPair.second;
		if(methodTable->hasErrors()) {
			foreach(std::string error, methodTable->getErrors()) errors.push_back(error);
		}
	}

	visit(program);
	if(variables.hasErrors()) {
		foreach(std::string error, variables.getErrors()) errors.push_back(error);
	}
}

void TypeChecker::visit(shared_ptr<Program> program) {
	visit(program->mainClass);

	foreach(shared_ptr<ClassDecl> decl, program->classList->classDecls) visit(decl);
}

void TypeChecker::visit(shared_ptr<MainClass> mainClass) {
	variables.beginScope();
	classScope = mainClass->classId->name;

	foreach(shared_ptr<Stmt> stmt, mainClass->stmtList->stmts) visit(stmt);

	classScope = "";
	variables.endScope();
}

void TypeChecker::visit(shared_ptr<ClassDecl> classDecl) {
	switch(classDecl->getType()) {
		case NormalClassType:
			visit(shared_static_cast<NormalClassDecl>(classDecl));
		case ExtendsClassType:
			visit(shared_static_cast<ExtendsClassDecl>(classDecl));
	}
}

void TypeChecker::visit(shared_ptr<NormalClassDecl> normalClassDecl) {
	variables.beginScope();
	classScope = normalClassDecl->classId->name;

	foreach(shared_ptr<VarDecl> decl, normalClassDecl->varList->varDecls) visit(decl);
	foreach(shared_ptr<MethodDecl> decl, normalClassDecl->methodList->methodDecls) visit(decl);

	classScope = "";
	variables.endScope();
}

void TypeChecker::visit(shared_ptr<ExtendsClassDecl> extendsClassDecl) {
	variables.beginScope();
	classScope = extendsClassDecl->classId->name;

	foreach(shared_ptr<VarDecl> decl, extendsClassDecl->varList->varDecls) visit(decl);
	foreach(shared_ptr<MethodDecl> decl, extendsClassDecl->methodList->methodDecls) visit(decl);

	classScope = "";
	variables.endScope();
}

void TypeChecker::visit(shared_ptr<VarDecl> varDecl) {
	std::string name = varDecl->id->name;
	std::string type = varDecl->type->name;
	shared_ptr<Symbol> newSymbol(new Symbol(name, VariableSymbol));
	newSymbol->setReturn(type);

	variables.addSymbol(newSymbol);
}

void TypeChecker::visit(shared_ptr<MethodDecl> methodDecl) {
	std::string name = methodDecl->id->name;
	std::string type = methodDecl->type->name;

	variables.beginScope();

	foreach(shared_ptr<Argument> argument, methodDecl->argumentList->arguments) visit(argument);
	foreach(shared_ptr<Stmt> stmt, methodDecl->stmtList->stmts) visit(stmt);
	if(visit(methodDecl->exp).name != type) errors.push_back("Return type expected was " + type);

	variables.endScope();
}

void TypeChecker::visit(shared_ptr<Argument> argument) {
	std::string name = argument->id->name;
	std::string type = argument->type->name;
	shared_ptr<Symbol> newSymbol(new Symbol(name, VariableSymbol));
	newSymbol->setReturn(type);

	variables.addSymbol(newSymbol);
}

void TypeChecker::visit(shared_ptr<Stmt> stmt) {
	switch(stmt->getType()) {
		case NestedType:
			return visit(shared_static_cast<NestedStmt>(stmt));
		case IfType:
			return visit(shared_static_cast<IfStmt>(stmt));
		case WhileType:
			return visit(shared_static_cast<WhileStmt>(stmt));
		case PrintlnType:
			return visit(shared_static_cast<PrintlnStmt>(stmt));
		case AssignmentType:
			return visit(shared_static_cast<AssignmentStmt>(stmt));
		case ArrayAssignmentType:
			return visit(shared_static_cast<ArrayAssignmentStmt>(stmt));
		case VarDeclType:
			return visit(shared_static_cast<VarDeclStmt>(stmt));
	}
}

void TypeChecker::visit(shared_ptr<NestedStmt> nestedStmt) {
	variables.beginScope();
	foreach(shared_ptr<Stmt> stmt, nestedStmt->stmtList->stmts) visit(stmt);	
	variables.endScope();
}

void TypeChecker::visit(shared_ptr<IfStmt> ifStmt) {
	if(visit(ifStmt->exp).name != "bool") errors.push_back("Expected boolean expression");
	visit(ifStmt->ifStmt);
	visit(ifStmt->elseStmt);
}

void TypeChecker::visit(shared_ptr<WhileStmt> whileStmt) {
	if(visit(whileStmt->exp).name != "bool") errors.push_back("Expected boolean expression");
	visit(whileStmt->stmt);
}

void TypeChecker::visit(shared_ptr<PrintlnStmt> printlnStmt) {
	if(visit(printlnStmt->exp).name != "int") errors.push_back("Expected integer expression in System.out.println.");
}

void TypeChecker::visit(shared_ptr<AssignmentStmt> assignmentStmt) {
	shared_ptr<Symbol> symbol = variables.getSymbol(assignmentStmt->id->name);
	if(symbol->getType() == UndefinedSymbol) {
		errors.push_back("Undefined symbol " + symbol->getName());
		return;
	}

	if(symbol->getReturn() != visit(assignmentStmt->exp).name) errors.push_back("Type in assignment wrong.");
}

void TypeChecker::visit(shared_ptr<ArrayAssignmentStmt> arrayAssignmentStmt) {
	shared_ptr<Symbol> symbol = variables.getSymbol(arrayAssignmentStmt->id->name);
	if(symbol->getType() == UndefinedSymbol) {
		errors.push_back("Undefined symbol " + symbol->getName());
		return;
	}

	if(symbol->getReturn() != "int[]") errors.push_back("Variable " + arrayAssignmentStmt->id->name + " not an array");
	if(visit(arrayAssignmentStmt->exp).name != "int") errors.push_back("Expected integer expression.");
	if(visit(arrayAssignmentStmt->assignmentExp).name == symbol->getReturn()) errors.push_back("Expected integer expression.");
}

void TypeChecker::visit(shared_ptr<VarDeclStmt> varDeclStmt) {
	visit(varDeclStmt->varDecl);
}

Type TypeChecker::visit(shared_ptr<Exp> exp) {
	switch(exp->getType()) {
		case OrType:
			return visit(shared_static_cast<OrExp>(exp));
		case AndType:
			return visit(shared_static_cast<AndExp>(exp));
		case LessThanType:
			return visit(shared_static_cast<LessThanExp>(exp));
		case GreaterThanType:
			return visit(shared_static_cast<GreaterThanExp>(exp));
		case DifferentType:
			return visit(shared_static_cast<DifferentExp>(exp));
		case EqualType:
			return visit(shared_static_cast<EqualExp>(exp));
		case PlusType:
			return visit(shared_static_cast<PlusExp>(exp));
		case MinusType:
			return visit(shared_static_cast<MinusExp>(exp));
		case TimesType:
			return visit(shared_static_cast<TimesExp>(exp));
		case ArrayValueType:
			return visit(shared_static_cast<ArrayValueExp>(exp));
		case ArrayLengthType:
			return visit(shared_static_cast<ArrayLengthExp>(exp));
		case MemberFunctionType:
			return visit(shared_static_cast<MemberFunctionExp>(exp));
		case NewIntArrayType:
			return visit(shared_static_cast<NewIntArrayExp>(exp));
		case NewIdType:
			return visit(shared_static_cast<NewIdExp>(exp));
		case NotType:
			return visit(shared_static_cast<NotExp>(exp));
		case NumberType:
			return visit(shared_static_cast<NumberExp>(exp));
		case TrueType:
			return visit(shared_static_cast<TrueExp>(exp));
		case FalseType:
			return visit(shared_static_cast<FalseExp>(exp));
		case IdType:
			return visit(shared_static_cast<IdExp>(exp));
		case ThisType:
			return visit(shared_static_cast<ThisExp>(exp));
	}
	return Type("_internal_error1");
}

Type TypeChecker::visit(shared_ptr<OrExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "bool")) errors.push_back("Expected boolean expressions.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<AndExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "bool")) errors.push_back("Expected boolean expressions.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<LessThanExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "bool")) errors.push_back("Expected boolean expressions.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<GreaterThanExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "bool")) errors.push_back("Expected boolean expressions.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<DifferentExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && (visit(exp->exp1).name == "bool" || visit(exp->exp1).name == "int"))) errors.push_back("Expected expressions with same type.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<EqualExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && (visit(exp->exp1).name == "bool" || visit(exp->exp1).name == "int"))) errors.push_back("Expected expressions with same type.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<PlusExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "int")) errors.push_back("Expected integer expressions.");

	return Type("int");
}

Type TypeChecker::visit(shared_ptr<MinusExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "int")) errors.push_back("Expected integer expressions.");

	return Type("int");
}

Type TypeChecker::visit(shared_ptr<TimesExp> exp) {
	if(!(visit(exp->exp1).name == visit(exp->exp2).name && visit(exp->exp1).name == "int")) errors.push_back("Expected integer expressions.");

	return Type("int");
}

Type TypeChecker::visit(shared_ptr<ArrayValueExp> exp) {
	return Type("int");
}

Type TypeChecker::visit(shared_ptr<ArrayLengthExp> exp) {
	return Type("int");
}

Type TypeChecker::visit(shared_ptr<MemberFunctionExp> exp) {
	Type classType = visit(exp->exp);
	shared_ptr<Symbol> classSymbol = classTable->getSymbol(classType.name);
	if(classSymbol->getType() == UndefinedSymbol) {
		errors.push_back("Class " + classType.name + " undeclared");
		return Type(classType.name);
	}

	shared_ptr<SymbolTable> classMethods = methods[classType.name];
	shared_ptr<Symbol> methodSymbol = classMethods->getSymbol(exp->id->name);
	if(methodSymbol->getType() == UndefinedSymbol) {
		errors.push_back("Method " + exp->id->name + " doesn't exist in class " + classType.name);
		return Type("_internal_error2");
	}

	std::vector<std::string> arguments = methodSymbol->getArguments();
	std::vector<shared_ptr<Exp> > list = exp->list->exps;

	if(list.size() != arguments.size()) {
		errors.push_back("Argument number wrong in method " + methodSymbol->getName() + " in class " + classType.name);
		return Type(methodSymbol->getReturn());
	}

	for(unsigned int i = 0; i < list.size(); i++) {
		if(visit(list[i]).name != arguments[i]) {
			errors.push_back("Expected argument with type " + arguments[i] + " in method " + methodSymbol->getName() + " in class " + classType.name);
		}
	}

	return Type(methodSymbol->getReturn());
}

Type TypeChecker::visit(shared_ptr<NewIntArrayExp> exp) {
	return Type("int[]");
}

Type TypeChecker::visit(shared_ptr<NewIdExp> exp) {
	return Type(exp->id->name);
}

Type TypeChecker::visit(shared_ptr<NotExp> exp) {
	shared_ptr<Exp> expression = exp->exp;	
	if(visit(expression).name != "bool") errors.push_back("Expected expression to be boolean.");

	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<NumberExp> exp) {
	return Type("int");
}

Type TypeChecker::visit(shared_ptr<TrueExp> exp) {
	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<FalseExp> exp) {
	return Type("bool");
}

Type TypeChecker::visit(shared_ptr<IdExp> exp) {
	shared_ptr<Symbol> symbol = variables.getSymbol(exp->id->name);
	if(symbol->getType() == UndefinedSymbol) {
		errors.push_back("Undeclared variable with name " + exp->id->name);
	}

	return Type(symbol->getReturn());
}

Type TypeChecker::visit(shared_ptr<ThisExp> exp) {
	return Type(classScope);
}

bool TypeChecker::hasErrors() {
	return (!errors.empty());
}

std::vector<std::string> TypeChecker::getErrors() {
	return errors;
}
